#include <Arduino.h>

#include <SPI.h>
#include <avr/pgmspace.h>

#include "long_flashData.h"
#include "long_eepData.h"

// #include "short_flashData.h"
// #include "short_eepData.h"

// Arduino UNO    ATTiny44A
// ----------------------------
// D10 (RESET)    RESET (Pin 4)
// D11 (MOSI)     MOSI (Pin 5)
// D12 (MISO)     MISO (Pin 6)
// D13 (SCK)      SCK  (Pin 7)
// GND            GND  (Pin 3)
// 5V             VCC  (Pin 1)
#define PIN_NRST 10
#define PIN_MOSI 11
#define PIN_MISO 12
#define PIN_SCK  13
#define PIN_KEY_CALIBRATION  2
#define PIN_KEY_PROG   3
#define PIN_RED_LED    4
#define PIN_GREEN_LED  5

void leaveProgrammingMode();
void enterProgrammingMode();
bool checkSignature();
void programFlash(const uint8_t *data, uint16_t size);
uint8_t sendProgrammingCommand(uint8_t cmd1, uint8_t cmd2, uint8_t cmd3, uint8_t cmd4);
void eraseChip();
bool readFlash(uint16_t startAddr, uint16_t length);
void programEEPROM(const uint8_t *data, uint16_t size);
bool readEEPROM(uint16_t startAddr, uint16_t length);

// 定义 ATTiny44A 的签名
#define ATTINY44A_SIGNATURE 0x1E9207

// 命令定义
#define PROG_COMMAND_ENABLE_PROG    0xAC, 0x53
#define PROG_COMMAND_READ_SIGNATURE 0x30
#define PROG_COMMAND_READ_FLASH     0x20
#define PROG_COMMAND_READ_EEPROM    0xA0
#define PROG_COMMAND_WRITE_FLASH    0x40
#define PROG_COMMAND_WRITE_EEPROM   0xC0
#define PROG_COMMAND_ERASE_CHIP     0xAC, 0x80
#define PROG_COMMAND_WRITE_FUSE     0xAC
#define PROG_COMMAND_READ_FUSE      0x50
#define PROG_COMMAND_READ_LOCK      0x58
#define PROG_COMMAND_WRITE_LOCK     0xAC


// 进入编程模式
#define SPI_CLOCK 		(1000000/6)
void enterProgrammingMode() {
  pinMode(PIN_NRST, OUTPUT);
  digitalWrite(PIN_NRST, HIGH);
  SPI.begin();
  SPI.beginTransaction(SPISettings(SPI_CLOCK, MSBFIRST, SPI_MODE0));
  digitalWrite(PIN_SCK, LOW);
  digitalWrite(PIN_NRST, LOW);
  delay(20);
  sendProgrammingCommand(PROG_COMMAND_ENABLE_PROG, 0x00, 0x00);
}

// 退出编程模式
void leaveProgrammingMode() {
  SPI.endTransaction();
  digitalWrite(10, HIGH);
}

// 发送编程命令
uint8_t sendProgrammingCommand(uint8_t cmd1, uint8_t cmd2, uint8_t cmd3, uint8_t cmd4) {
  SPI.transfer(cmd1);
  SPI.transfer(cmd2);
  SPI.transfer(cmd3);
  return SPI.transfer(cmd4);
}

// 检查设备签名
bool checkSignature() {
  uint32_t sig1 = sendProgrammingCommand(PROG_COMMAND_READ_SIGNATURE, 0x00, 0x00, 0x00);
  uint32_t sig2 = sendProgrammingCommand(PROG_COMMAND_READ_SIGNATURE, 0x00, 0x01, 0x00);
  uint32_t sig3 = sendProgrammingCommand(PROG_COMMAND_READ_SIGNATURE, 0x00, 0x02, 0x00);

  uint32_t signature = (sig1 << 16) | (sig2 << 8) | sig3;

  Serial.print("Device signature: 0x");
  Serial.println(signature, HEX);

  return (signature == ATTINY44A_SIGNATURE);
}

// 擦除芯片
void eraseChip() {
  Serial.println("Erasing chip...");
  sendProgrammingCommand(PROG_COMMAND_ERASE_CHIP, 0x00, 0x00);
  delay(100); // 等待擦除完成
  Serial.println("Chip erased");
}

// 烧录Flash
void programFlash(const uint8_t *data, uint16_t size) {
  Serial.println("Programming Flash...");

  for (uint16_t addr = 0; addr < size; addr += 2) {
    // 加载低字节
    uint8_t lowByte = pgm_read_byte(&data[addr]);
    // 加载高字节 (如果存在)
    uint8_t highByte = (addr + 1 < size) ? pgm_read_byte(&data[addr + 1]) : 0xFF;

    // 写入低字节
    sendProgrammingCommand(PROG_COMMAND_WRITE_FLASH, (addr >> 9), addr >> 1, lowByte);
    delay(1);

    // 写入高字节
    sendProgrammingCommand(PROG_COMMAND_WRITE_FLASH | 0x08, (addr >> 9), addr >> 1, highByte);
    delay(1);

    // 编程Flash页 (每页64字节)
    if ((addr % 64 == 62) || (addr == size - 2)) {
      sendProgrammingCommand(0x4C, (addr >> 9), addr >> 1, 0x00);
      delay(20);
    }

    if (addr % 64 == 0) {
      Serial.print(".");
    }
  }

  Serial.println("\nFlash programming complete");
}

// 烧录EEPROM
void programEEPROM(const uint8_t *data, uint16_t size) {
  Serial.println("Programming EEPROM...");

  for (uint16_t addr = 0; addr < size; addr++) {
    uint8_t byteValue = pgm_read_byte(&data[addr]);

    sendProgrammingCommand(PROG_COMMAND_WRITE_EEPROM, 0, addr, byteValue);
    delay(10);

    if (addr % 16 == 0) {
      Serial.print(".");
    }
  }

  Serial.println("\nEEPROM programming complete");
}

// 读取Flash
bool readFlash(uint16_t startAddr, uint16_t length) {
  Serial.println("Reading Flash...");
  Serial.print("Address  | Data");
  Serial.println("----------------");

  for (uint16_t addr = startAddr; addr < startAddr + length; addr += 2) {
    // 读取低字节
    uint8_t lowByte = sendProgrammingCommand(PROG_COMMAND_READ_FLASH, (addr >> 9), addr >> 1, 0x00);
    // 读取高字节
    uint8_t highByte = sendProgrammingCommand(PROG_COMMAND_READ_FLASH | 0x08, (addr >> 9), addr >> 1, 0x00);

    uint8_t lowByte_t = pgm_read_byte(&flashData[addr]);
    uint8_t highByte_t = pgm_read_byte(&flashData[addr+1]);

    if(lowByte_t != lowByte || highByte_t != highByte){
      Serial.print("Error: Flash data does not match: ");
      Serial.print("0x");
      Serial.print(addr, HEX);
      Serial.print(" | 0x");
      if (lowByte < 0x10) Serial.print("0");
      Serial.print(lowByte, HEX);
      Serial.print(" 0x");
      if (highByte < 0x10) Serial.print("0");
      Serial.println(highByte, HEX);
      return false;
    }
    if(addr % 128 == 0){
      Serial.print(".");
      digitalWrite(PIN_GREEN_LED, !digitalRead(PIN_GREEN_LED));
    }
  }
  Serial.println("Flash read complete: All data matches");
  digitalWrite(PIN_GREEN_LED, 0);
  return true;
}

// 读取EEPROM
bool readEEPROM(uint16_t startAddr, uint16_t length) {
  Serial.println("Reading EEPROM...");
  Serial.print("Address | Data");
  Serial.println("---------------");

  for (uint16_t addr = startAddr; addr < startAddr + length; addr++) {
    uint8_t byteValue = sendProgrammingCommand(PROG_COMMAND_READ_EEPROM, 0, addr, 0x00);
    uint8_t byteValue_t = pgm_read_byte(&eepData[addr]);

    if(byteValue_t != byteValue){
      Serial.print("Error: Eeprom data does not match: ");
      Serial.print("0x");
      Serial.print(addr, HEX);
      Serial.print(" | 0x");
      if (byteValue < 0x10) Serial.print("0");
      Serial.println(byteValue, HEX);
      return false;
    }
    if(addr % 32 == 0){
      Serial.print(".");
    }
  }
  Serial.println("Eeprom read complete: All data matches");
  return true;
}

uint8_t tim100msCnt = 0;
uint8_t tim500msCnt = 0;
uint32_t last_now = 0;
uint8_t keyCalibCnt = 0;
bool keyCalibFlag = false;
bool keyCalibCheck = false;
uint8_t keyProgCnt = 0;
bool keyProgFlag = false;
bool keyProgCheck = false;

uint8_t CalibStep = 0;
uint8_t ProgStep = 0;

void setup() {
  Serial.begin(9600);
  pinMode(PIN_KEY_CALIBRATION, INPUT_PULLUP);
  pinMode(PIN_KEY_PROG, INPUT_PULLUP);
  pinMode(PIN_RED_LED, OUTPUT);
  pinMode(PIN_GREEN_LED, OUTPUT);
  // 示例操作 - 可以根据需要取消注释
  // eraseChip(); // 擦除芯片
  // programFlash((const uint8_t *)flashData, flashData_len);
  // readFlash(0, flashData_len); // 读取前32字节Flash
  // programEEPROM((const uint8_t *)eepData, eepData_len); // 烧录EEPROM
  // readEEPROM(0, eepData_len); // 读取前16字节EEPROM
  last_now = micros();
  randomSeed(analogRead(0));
}

void loop() {
  if(keyCalibFlag){
    keyCalibFlag = false;
    enterProgrammingMode();
    if (checkSignature()) {
      bool result = readFlash(0, flashData_len); // 读取前32字节Flash
      if(result){
        digitalWrite(PIN_RED_LED, 1);
        digitalWrite(PIN_GREEN_LED, 0);
      } else {
        digitalWrite(PIN_RED_LED, 0);
        digitalWrite(PIN_GREEN_LED, 1);
        result = readEEPROM(0, eepData_len); // 读取前32字节Flash
        if(result){
          digitalWrite(PIN_RED_LED, 1);
          digitalWrite(PIN_GREEN_LED, 0);
        } else {
          digitalWrite(PIN_RED_LED, 0);
          digitalWrite(PIN_GREEN_LED, 1);
        }
      }
      if(!result){
        uint8_t randomNum = random(100);
        if(randomNum != 50){
          //prog
        }
      }
    } else {
      digitalWrite(PIN_RED_LED, 1);
      digitalWrite(PIN_GREEN_LED, 0);
    }
    leaveProgrammingMode();
  }

  if(keyProgFlag){
    keyProgFlag = false;
  }

  if(micros() - last_now > 10){
    last_now = micros();
    tim100msCnt++;
    tim500msCnt++;
    bool calibFlag = (digitalRead(PIN_KEY_CALIBRATION) == LOW);
    bool progFlag = (digitalRead(PIN_KEY_PROG) == LOW);
    if(calibFlag){
      if(keyCalibCnt < 20) keyCalibCnt++;
    } else {
      if(keyCalibCnt >= 20){ // 按键短按200ms
        keyCalibFlag = true;
        keyCalibCheck = false;
      } else if((keyCalibCnt > 10) && !keyCalibCheck){
        keyCalibCheck = true;
      }

      if(keyCalibCheck) keyCalibCnt--;

      if((keyCalibCnt < 3) || !keyCalibCheck){
        keyCalibCheck = false;
        keyCalibCnt = 0;
      }
    }

    if(progFlag){
      if(keyProgCnt < 20) keyProgCnt++;
    } else {
      if(keyProgCnt >= 20){ // 按键短按200ms
        keyProgFlag = true;
        keyProgCheck = false;
      } else if((keyProgCnt > 10) && !keyProgCheck){
        keyProgCheck = true;
      }

      if(keyProgCheck) keyProgCnt--;

      if((keyProgCnt < 3) || !keyProgCheck){
        keyProgCheck = false;
        keyProgCnt = 0;
      }
    }
    
    if(tim100msCnt >= 10){
      tim100msCnt = 0;
    }
    if(tim500msCnt >= 50){
      tim500msCnt = 0;
    }
  }
}